﻿using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;

namespace Affinity.Core.Collections
{
    /// <summary>
    /// A class providing a modifiable implementation of ILookup.
    /// </summary>
    /// <typeparam name="TKey"></typeparam>
    /// <typeparam name="TValue"></typeparam>
    public class ModifiableLookup<TKey, TValue> : ILookup<TKey, TValue>
    {
        /// <summary>
        /// Adds a value for a key.
        /// </summary>
        /// <param name="key"></param>
        /// <param name="value"></param>
        public void Add(TKey key, TValue value)
        {
            if (!_KeyToValuesMap.ContainsKey(key)) _KeyToValuesMap[key] = new List<TValue>();
            _KeyToValuesMap[key].Add(value);
        }

        /// <summary>
        /// Adds values for a key.
        /// </summary>
        /// <param name="key"></param>
        /// <param name="list"></param>
        public void Add(TKey key, IEnumerable<TValue> list)
        {
            if (!_KeyToValuesMap.ContainsKey(key)) _KeyToValuesMap[key] = new List<TValue>();
            foreach (var item in list)
            {
                _KeyToValuesMap[key].Add(item);
            }
        }

        /// <summary>
        /// Returns whether a key is contained in this lookup.
        /// </summary>
        /// <param name="key"></param>
        /// <returns></returns>
        public bool Contains(TKey key)
        {
            return _KeyToValuesMap.ContainsKey(key);
        }

        /// <summary>
        /// Returns the number of keys.
        /// </summary>
        public int Count
        {
            get { return _KeyToValuesMap.Count; }
        }

        /// <summary>
        /// Gets an enumeration of values for a key.
        /// </summary>
        /// <param name="key"></param>
        /// <returns></returns>
        public IEnumerable<TValue> this[TKey key]
        {
            get { return _KeyToValuesMap.ContainsKey(key) ? (IEnumerable<TValue>)_KeyToValuesMap[key] : new TValue[0]; }
        }

        /// <summary>
        /// Gets an enumerator over keyed groupings.
        /// </summary>
        /// <returns></returns>
        public IEnumerator<IGrouping<TKey, TValue>> GetEnumerator()
        {
            return _KeyToValuesMap.Select(pair => new Grouping(pair.Key, pair.Value.ToArray())).Cast<IGrouping<TKey, TValue>>().GetEnumerator();
        }

        /// <summary>
        /// Removes a value from the collection corresponding to the given key.
        /// </summary>
        /// <param name="key"></param>
        /// <param name="value"></param>
        public bool Remove(TKey key, TValue value)
        {
            bool removed = false;
            if (_KeyToValuesMap.ContainsKey(key))
            {
                removed = _KeyToValuesMap[key].Remove(value);
                if (_KeyToValuesMap[key].Count == 0) _KeyToValuesMap.Remove(key);
            }
            return removed;
        }

        /// <summary>
        /// Removes values from the collection corresponding to the given key.
        /// </summary>
        /// <param name="key"></param>
        /// <param name="values"></param>
        /// <returns></returns>
        public void Remove(TKey key, IEnumerable<TValue> values)
        {
            if (values == null) throw new ArgumentNullException("values");
            foreach (TValue value in values) Remove(key, value);
        }

        /// <summary>
        /// Gets an enumerator over keyed groupings.
        /// </summary>
        /// <returns></returns>
        IEnumerator IEnumerable.GetEnumerator()
        {
            return GetEnumerator();
        }

        /// <summary>
        /// Represents a grouping by key.
        /// </summary>
        private class Grouping : IGrouping<TKey, TValue>
        {
            /// <summary>
            /// Constructor accepting the key and elements for this grouping.
            /// </summary>
            /// <param name="key"></param>
            /// <param name="elements"></param>
            public Grouping(TKey key, IEnumerable<TValue> elements)
            {
                _elements = elements;
                _key = key;
            }

            /// <summary>
            /// Gets the key for this grouping.
            /// </summary>
            public TKey Key
            {
                get { return _key; }
            }

            /// <summary>
            /// Gets an enumerator over the values in this grouping.
            /// </summary>
            /// <returns></returns>
            public IEnumerator<TValue> GetEnumerator()
            {
                return _elements.GetEnumerator();
            }

            /// <summary>
            /// Gest an enumerator over the values in this grouping.
            /// </summary>
            /// <returns></returns>
            IEnumerator IEnumerable.GetEnumerator()
            {
                return GetEnumerator();
            }

            private readonly TKey _key;
            private readonly IEnumerable<TValue> _elements;
        }

        private Dictionary<TKey, List<TValue>> _KeyToValuesMap = new Dictionary<TKey, List<TValue>>();
    }
}
